import random

import ahocorasick


class generate_testset:
    def __init__(self):
        # 设置各类别句子模板
        self.is_symptom_templates = [
            '{0}和{1}有什么关系？', '{0}和{1}有啥关系啊', '{0}和{1}有没有关系？', '{0}和{1}有什么联系',
            '{1}和{0}有什么关系', '{1}和{0}有啥关系？', '{1}和{0}有没有关系？', '{1}和{0}有什么联系',
            '得了{0}会有{1}？', '得了{0}一般会引起{1}这种症状嘛？', '患上{0}的人会有{1}的症状咩',
            '{0}会不会造成{1}？', '得了{0}会有{1}的症状嘛？', '一个人如果得了{0}会有{1}吗？',
            '{1}是不是{0}的症状之一', '{1}是不是{0}的症状', '{1}是{0}的症状之一吗'
        ]

        self.is_food_templates = [
            '得了{0}可以吃{1}吗？', '得了{0}可以食用{1}吗？', '患上{0}可以饮用{1}嘛', '得了{0}能够摄入{1}吗？', '患了{0}可以食用{1}咩？',
            '得了{0}的病人可以摄入{1}吗？', '患上{0}的病人可不可以吃{1}', '{1}对{0}的康复或缓解有影响吗？', '食用{1}对得了{0}的病人有好处吗',
            '摄入{1}对{0}的康复有益吗', '食用{1}对{0}的康复有害吗', '{1}对{0}的缓解有不良影响嘛？', '{1}是{0}患者的忌食食物吗？',
            '{1}是{0}患者的宜吃食物吗？', '{1}是{0}病人的宜吃食品吗', '{1}适合患了{0}的人吃吗？'
        ]

        self.is_belonging_templates = [
            '{1}可以治疗{0}吗？', '{1}能够治疗{0}嘛', '{1}能不能治疗{0}', '{1}可不可以治{0}', '得了{0}可以到{1}去治疗吗？',
            '得了{0}可以去{1}挂号看病吗？', '患上{0}一般是到{1}看医生吗', '我可以去{1}挂号看{0}吗', '如果我得了{0}，是不是到{1}去挂号看病？',
            '假如一个人得了{0}，可以到{1}看吗', '{0}和{1}是啥关系', '{0}和{1}有什么联系', '{1}和{0}有什么关系'
        ]

        self.is_recommend_drug_templates = [
            '{1}是不是治{0}的药物', '{1}可以治疗{0}吗？', '{1}能不能治{0}？', '{0}的推荐药品是{1}嘛', '{1}是{0}的推荐药吗？',
            '得了{0}推荐吃{1}吗', '得了{0}可以吃{1}咩？', '{0}是吃{1}吗', '一般{0}推荐吃{1}吗', '如果我得了{0}，是不是应该吃{1}治疗',
            '{1}对{0}有治疗效果嘛？'
        ]

        self.is_need_check_templates = [
            '{1}可以诊断出{0}吗', '{1}能够查出{0}吗？', '{1}能够试出{0}吗', '一般做{1}是否能查出{0}？', '{1}可以检测出{0}嘛？',
            '{1}是{0}的检查项目吗', '{1}可不可以检查出{0}', '我得了{0}，可以通过{1}检测出来吗', '得了{0}要做{1}吗'
        ]

        self.others_templates = [
            '{0}有啥症状？', '{0}有什么症状？', '{0}的症状是什么？', '什么是{0}', '{0}是什么？', '{0}有什么并发症吗？',
            '{0}的并发症是啥', '{0}的病因一般是啥？', '{0}的宜吃食物是什么？', '得了{0}建议吃点啥？', '{0}的忌食食品有啥？', '得了{0}的人不能吃什么？', '{0}病人最好不要吃什么？',
            '吃什么食品对患了{0}的病人有好处', '得了{0}一般要吃啥药？', '{0}患者一般要做什么检查', '{0}做什么检查项目能够查出来？', '{0}有什么预防措施',
            '{0}的治愈率是多少？', '{0}一般会引起什么症状？', '{0}的目标人群？', '{0}有什么推荐药物吗', '得了{0}一般到哪个科室挂号？', '得了{0}一般到哪里看？',
            '{0}的易感人群是哪些人？'
        ]

        self.others_templates_food = [
            '{0}对治疗什么病有好处', '{0}是啥病的推荐食物？', '得了什么病适合吃{0}', '{0}对什么病的康复有阻碍', '啥病推荐吃{0}'
        ]

        self.other_templates_check = [
            '{0}一般能查出什么病？', '得了什么病可以通过{0}测出来', '{0}可以检查出啥病症？'
        ]

        self.other_templates_department = [
            '{0}一般治疗那些疾病？', '一般到{0}看啥病？', '什么病人是去{0}挂号？'
        ]

        self.other_templates_drug = [
            '{0}是啥病的推荐药物？', '{0}是什么病的推荐药品', '{0}可以治疗什么病？', '得了什么病推荐吃{0}'
        ]

        self.other_templates_symptom = [
            '{0}的症状一般在什么病当中会出现？', '得了什么病会有{0}的症状', '{0}是啥病的症状？', '什么病会出现{0}'
        ]

    def filter_repeated_wds(self, sentences, actree, words):
        """
        只调用一次，用于生成训练集和验证集中都没有覆盖的元素
        @param sentences: 问题句子列表
        @param actree: 领域类AC自动机，用于加速匹配
        @param words: 领域词汇表
        @return: 滤除重复词的领域词表
        """
        repeated_wds = set()
        for question in sentences:
            for i in actree.iter(question):  # .iter() 子串匹配 i:[end_index, [insert_order, original_value]]
                wd = i[1][1]
                repeated_wds.add(wd)

        return list(set(words) - repeated_wds)

    def load_data(self, path, is_dataset=False):
        """
        数据加载
        @param path: 文件路径
        @return: 句子列表
        """
        data = []
        file = open(path, 'r', encoding='utf8', errors='ignore')
        file_data = file.readlines()  # 读取所有行
        if is_dataset:
            for row in file_data[1:]:
                if row == '\n' or len(row) == 0:
                    continue
                row = row.replace('\n', '')
                tmp_list = row.split('\t')
                data.append(tmp_list[1])
        else:
            for row in file_data:
                row = row.replace('\n', '')
                data.append(row)
        print('------ {0} loaded... ------'.format(path.split('\\')[-1]))
        return data

    def build_actree(self, wordlist):
        actree = ahocorasick.Automaton()
        for index, word in enumerate(wordlist):
            actree.add_word(word, (index, word))
        actree.make_automaton()
        return actree

    def generate_sentence(self, label, templates, entities_1, entities_2, each_num):
        """
        生成对应label的问句
        @param label: 问题类别标签
        @param templates: 问句模板
        @param entities_1: 填充用实体列表
        @param entities_2: 填充用实体列表
        @param each_num: 每个模板生成的问句数目
        @return: 字典列表。[{'label':'', 'text': ''}]
        """
        result = []
        for template in templates:
            for i in range(each_num):
                A = random.choice(entities_1)

                if entities_2 != []:
                    B = random.choice(entities_2)
                    text = template.format(A, B)
                else:
                    text = template.format(A)

                result.append({'label': label, 'text': text})

        return result

    def save_data(self, data, path):
        """
        保存数据
        @param data: 保存数据
        @param path: 保存路径
        @return: Boolean
        """
        fileObject = open(path, 'w')
        fileObject.write('label\ttext\n')
        for dictionary in data:
            fileObject.write(dictionary['label'] + '\t' + dictionary['text'] + '\n')
        fileObject.close()


if __name__ == '__main__':
    test_gen = generate_testset()

    train = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\data\dataset_ernie_tiny\train.txt',
        True)
    dev = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\data\dataset_ernie_tiny\dev.txt', True)

    # 读取领域词
    check_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\check.txt')
    department_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\department.txt')
    disease_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\disease.txt')
    drug_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\drug.txt')
    food_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\food.txt')
    symptom_wds = test_gen.load_data(
        r'C:\CIKE\diploma project\Git\Back\kg-search-backend\QASystemOnMedicalKG\dict\symptom.txt')

    # 构造Actree加速匹配
    check_tree = test_gen.build_actree(check_wds)
    department_tree = test_gen.build_actree(department_wds)
    disease_tree = test_gen.build_actree(disease_wds)
    drug_tree = test_gen.build_actree(drug_wds)
    food_tree = test_gen.build_actree(food_wds)
    symptom_tree = test_gen.build_actree(symptom_wds)

    # 过滤
    sentences = train + dev
    clean_check_wds = test_gen.filter_repeated_wds(sentences, check_tree, check_wds)
    clean_department_wds = test_gen.filter_repeated_wds(sentences, department_tree, department_wds)
    clean_disease_wds = test_gen.filter_repeated_wds(sentences, disease_tree, disease_wds)
    clean_drug_wds = test_gen.filter_repeated_wds(sentences, drug_tree, drug_wds)
    clean_food_wds = test_gen.filter_repeated_wds(sentences, food_tree, food_wds)
    clean_symptom_wds = test_gen.filter_repeated_wds(sentences, symptom_tree, symptom_wds)
    print('------ filter finished ------')
    # 生成句子
    result_is_symptom = test_gen.generate_sentence('is_symptom', test_gen.is_symptom_templates, clean_disease_wds,
                                                   clean_symptom_wds, 5)
    result_is_food = test_gen.generate_sentence('is_food', test_gen.is_food_templates, clean_disease_wds,
                                                clean_food_wds, 5)
    result_is_belonging = test_gen.generate_sentence('is_belonging', test_gen.is_belonging_templates, clean_disease_wds,
                                                     clean_department_wds, 5)
    result_is_need_check = test_gen.generate_sentence('is_need_check', test_gen.is_need_check_templates,
                                                      clean_disease_wds, clean_check_wds, 5)
    result_is_recommend_drug = test_gen.generate_sentence('is_recommend_drug', test_gen.is_recommend_drug_templates,
                                                          clean_disease_wds, clean_drug_wds, 5)
    result_other = test_gen.generate_sentence('other', test_gen.others_templates, clean_disease_wds, [],
                                              5) + test_gen.generate_sentence('other', test_gen.other_templates_check,
                                                                              clean_check_wds, [],
                                                                              5) + test_gen.generate_sentence('other',
                                                                                                              test_gen.other_templates_department,
                                                                                                              clean_department_wds,
                                                                                                              [],
                                                                                                              5) + test_gen.generate_sentence(
        'other', test_gen.other_templates_drug,
        clean_drug_wds, [], 5) + test_gen.generate_sentence('other', test_gen.other_templates_symptom,
                                                            clean_symptom_wds, [], 5) + test_gen.generate_sentence(
        'other', test_gen.others_templates_food,
        clean_food_wds, [], 5)
    print('------ sentences generated ------')
    result_all = result_is_symptom + result_is_food + result_is_belonging + result_is_need_check + result_is_recommend_drug + result_other
    print('------ Total Length: {0} ------'.format(len(result_all)))

    # 保存文件
    random.shuffle(result_all)
    test_gen.save_data(result_all, './clean_test.txt')
